import ast
import json
from typing import Literal, Union, List

from asgiref.sync import sync_to_async
from langchain_chroma import Chroma
from langchain_core.documents import Document
from openai import OpenAI, AsyncOpenAI
from openai.types.beta import VectorStore

from llm.entity.dto.DeviceConfigAddDTO import DeviceConfigAddDTO
from llm.entity.metadatas import VectorStoreMetadata

from llm.entity.dao.VectorStoreDAO import VectorStoreDAO
from llm.entity.vo.DeviceConfigAddVO import DeviceConfigAddVO
from llm.models import VectorStores
from users.models import Users
from HomeAutoAI.common.utils.ThreadLocal import get_cur_user

client = OpenAI()


class VectorStoreService:
    @classmethod
    def create(cls, vector_name: str, assistant_id: str, user: Users,
               metadata: VectorStoreMetadata, file_ids: list[str] = None
               ) -> VectorStore:
        # TODO：metadata需要敲定
        if file_ids:
            vector_store = client.beta.vector_stores.create(
                file_ids=file_ids,
                name=vector_name,
                metadata=metadata.__dict__,
            )
        else:
            # 创建空向量库
            vector_store = client.beta.vector_stores.create(
                name=vector_name,
                metadata=metadata.__dict__,
            )
        client.beta.assistants.update(
            assistant_id=assistant_id,
            tool_resources={"file_search": {"vector_store_ids": [vector_store.id]}},
        )
        VectorStoreDAO.add(vector_store, user=user)
        return vector_store

    @classmethod
    def attach_file(cls, deviceConfigAddDTO: DeviceConfigAddDTO, user: Users,
                    vector_store_id: str = None, version: Literal["openai", "langgraph"] = "openai") -> DeviceConfigAddVO:
        """
        向向量库中添加文件,如果没给指定向量库，则默认添加到默认向量库中,
        否则就添加到指定的向量库中
        :param version: 指定程序的版本，初版使用的是”openai“的sdk和assistant框架;第二版使用的是Langgraph和Chroma数据库
        :param vector_store_id:
        :param deviceConfigAddDTO:
        :param user:
        :return:
        """
        vector_store = VectorStoreService._db_get_default(user)
        if version == "openai":
            # 添加给全屋助手
            # 获取当前用户的全屋助手id，以及默认的向量库id，进行对向量库的更新操作、

            if vector_store:
                # 将文件添加到已有的向量库中
                # TODO：待测试
                client.beta.vector_stores.file_batches.create_and_poll(
                    vector_store_id=vector_store.vector_store_id,
                    file_ids=deviceConfigAddDTO.file_ids,
                )
                vo = DeviceConfigAddVO()
                vo.file_ids = deviceConfigAddDTO.file_ids
                vo.device_id = deviceConfigAddDTO.device_id
                vo.vector_store_id = vector_store.vector_store_id
                return vo
            else:
                raise Exception("没找到默认向量库")
        elif version == "langgraph":
            if vector_store:
                vector_store = VectorStoreService._db_get_default(user)
                # 缓存区域的用户上传的文档存储在此数据库中
                cache_store = Chroma(
                collection_name="cache_area",
                persist_directory=r"D:\DevelopFiles\pycharms\HomeAutoAI\llm\ChromaDB\CacheArea",
                )
                # 用户对应的配置信息存储集合
                user_vector_store=Chroma(
                    collection_name=vector_store.vector_store_id,
                    persist_directory=r"D:\DevelopFiles\pycharms\HomeAutoAI\llm\ChromaDB\DeviceConfigs",
                )

                # 获取对应的文档内容，重新封装存入用户对应的集合区域中
                all=cache_store.get(
                    ids=deviceConfigAddDTO.file_ids
                )
                documents=[]
                for i in all["documents"]:
                    doc = Document(page_content=i)
                    documents.append(doc)

                user_vector_store.add_documents(documents)

                # 装配返回对象
                vo = DeviceConfigAddVO()
                vo.file_ids = deviceConfigAddDTO.file_ids
                vo.device_id = deviceConfigAddDTO.device_id
                vo.vector_store_id = vector_store.vector_store_id
                return vo
            else:
                raise Exception("没找到默认向量库")



    @classmethod
    def _db_get_default(cls, user: Users) -> VectorStores | None:
        # 这里数据库设计失误
        vector_stores:List[VectorStores] = VectorStores.objects.filter(user_id=user)
        print(f'vector_store_ids: {vector_stores}')
        for vector_store in vector_stores:
            # 对于这种带有单引号的字符串（例如 {'user_id': '', 'isInit': '', 'usage': ''}），
            # 它不符合 JSON 格式，无法直接使用 json.loads() 进行转换。可以使用 ast.literal_eval() 来安全地将其转换为字典。
            # ast.literal_eval() 可以处理 Python 字典表示的字符串。
            metadata = ast.literal_eval(vector_store.information)
            print(metadata)
            if metadata['isInit'] == 'True':
                return vector_store
        return None

